feat(pt/dpa4): add SO3 readout and small bug fix#5556
Conversation
📝 WalkthroughWalkthroughThree independent changes: (1) Changesso3_readout feature for DescrptSeZM
Deferred compile version check in SeZMModel
NvNeighborList CPU-only host improvements
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/utils/nv_nlist.py`:
- Around line 70-77: The current code structure risks leaking the file
descriptor saved_fd if either the open() call on devnull or the first dup2()
call raises an exception before reaching the existing try-finally block. To fix
this, wrap the os.dup(stderr_fd) call and the subsequent open() and try-finally
blocks with an outer try-finally statement that ensures os.close(saved_fd) is
called in the finally clause regardless of whether an exception occurs during
the devnull file opening or the initial dup2() redirection.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 313173f5-2344-4891-a24b-7c1d4e7883f7
📒 Files selected for processing (5)
deepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/utils/nv_nlist.pydeepmd/utils/argcheck.pyexamples/water/dpa4/input.json
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5556 +/- ##
==========================================
- Coverage 82.16% 80.85% -1.32%
==========================================
Files 896 896
Lines 102643 102676 +33
Branches 4340 4342 +2
==========================================
- Hits 84341 83020 -1321
- Misses 16965 18324 +1359
+ Partials 1337 1332 -5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Code reviewFound 1 blocking issue (CI is red because of it):
The DPA4/SeZM descriptor is not pt-only — deepmd-kit/deepmd/pt/model/descriptor/sezm.py Lines 2079 to 2081 in 1d244fd but Why this fails: There is also a quieter consequence for the non-default modes: dpmodel This is the serialization-discipline rule: a new serialized field must be added to all backends that share the format, in the same change — otherwise the shared-format consistency test breaks (as it has here). Suggested fix: mirror |
|
I may handle it after #5555 merges. |
## What Adds the DPA4/SeZM `so3_readout` option (`"none"` / `"glu"` / `"mlp"`) **across all backends** (pt + dpmodel + pt_expt), making it cross-backend consistent. Builds on **#5556** (@OutisLi): its pt `so3_readout` commit (`refactor(dpa4): output ffn`) is included here with original authorship preserved; this PR adds the missing **dpmodel** counterpart so the shared DPA4 serialize format round-trips across backends. (The unrelated `nv_nlist`/compiler-check fixes from #5556 are intentionally left to #5556.) ## Why `so3_readout` is implemented by configuring the final output FFN — `"glu"`/`"mlp"` turn on the SO(3)-grid FFN (`ffn_so3_grid`, `grid_mlp`). On its own (pt-only, as in #5556) it breaks DPA4 cross-backend consistency: pt `serialize()` emits `so3_readout` but dpmodel `DescrptDPA4` couldn't round-trip it → `source/tests/consistent/descriptor/test_dpa4.py::...::test_pt_consistent_with_ref` failed on every Test Python shard. This is now feasible and small because **#5555** already ported `ffn_so3_grid` + the SO(3)-grid machinery (`SO3GridNet`/`GridMLP`/`GridProduct`) into the dpmodel `EquivariantFFN`. So the dpmodel `so3_readout` is just: accept the param, configure `output_ffn` exactly like pt, wire the readout (l=0 slice for `"none"`; full `(N,D,1,C)` fold for `"glu"`/`"mlp"` then slice l=0), and serialize the key. pt_expt auto-wraps. ## Changes - **pt** (from #5556): `so3_readout` in `DescrptSeZM` + argcheck + `examples/water/dpa4/input.json`. - **dpmodel** `descriptor/dpa4.py`: `so3_readout` param + validation; `output_ffn` configured (`lmax=node_l_schedule[-1]`, `kmax=min(kmax, readout_lmax)`, `ffn_so3_grid`, `grid_mlp`, `grid_branch=0`); readout forward mirrors pt; serialize the key. - **pt_expt**: auto-wrapped (no explicit change). ## Validation - `test_dpa4.py` cross-backend consistency rows for `so3_readout ∈ {none, glu, mlp}` (pt vs dpmodel vs pt_expt, mixed_types) — green; `test_pt_consistent_with_ref` now passes. - Full-descriptor pt→dpmodel parity (weight-copied, `glu`+`mlp`) — **~7e-15 abs** (gate 1e-10), proving serialize interop. - dpa4 suite: 611 passed; ruff clean. ## Notes - Depends on #5555 (merged) for the dpmodel SO(3)-grid FFN. - `so3_readout` no longer "(Supported Backend: PyTorch)" — now multi-backend. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added a configurable `so3_readout` option to the DPA4 and SeZM descriptors (modes: `"none"`, `"glu"`, `"mlp"`), controlling how the final SO(3) readout is computed. * The setting is included in descriptor configuration serialization/deserialization to support round-tripping. * Updated the water DPA4 example to use `so3_readout: "mlp"`. * **Tests** * Added tests covering multiple readout modes, backend parity between implementations, and correct behavior for edge-free scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: OutisLi <LTC201806070316@gmail.com> Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Summary by CodeRabbit
New Features
so3_readoutconfiguration option for SeZM descriptor, enabling selection of computation methods ("none","glu","mlp") for final descriptor channel extraction.Improvements